#include <pmm.h>
#include <list.h>
#include <string.h>
#include <buddy_pmm.h>
#include <stdio.h>

#define LEFT_LEAF(index) ((index) * 2 + 1)
#define RIGHT_LEAF(index) ((index) * 2 + 2)
#define PARENT(index) ( ((index) + 1) / 2 - 1)
#define IS_POWER_OF_2(x) (!((x)&((x)-1)))
#define Max(a, b) ((a) > (b) ? (a) : (b))

uint8_t ROUND_DOWN_LOG(unsigned x) {
    int result = 0;
    while (x > 1) {
        x >>= 1;  // 使用右移操作将 x 除以 2
        result++;
    }
    return result;
}

free_area_t free_area;
#define nr_free (free_area.nr_free)


unsigned *buddy, buddy_size;
struct Page *base_page;

static void
buddy_init(void){
    nr_free = 0;
}

static void
buddy_init_memmap(struct Page *base, size_t n){
    assert(n > 0);
    struct Page *p = base;
    for (; p != base + n; p++)
    {
        assert(PageReserved(p));
        p->flags = p->property = 0;
        set_page_ref(p, 0);
        SetPageProperty(p);
    }
    base->property = n;
    SetPageProperty(base);
    nr_free += n;
    base_page = base;
    buddy_size =1<<ROUND_DOWN_LOG(n); 
    
    unsigned node_size = 2*buddy_size;  //内部节点数量
    buddy = (unsigned *)(base + buddy_size);
    //根据层数初始化节点大小
    for (int i = 0; i < 2 * buddy_size - 1; ++i)
    {
        if (IS_POWER_OF_2(i + 1))
            node_size >>= 1;
        buddy[i] = node_size;
    }
    
}

static struct Page *
buddy_alloc_pages(size_t n)
{
    assert(n > 0);
    if (n > nr_free)
    {
        return NULL;
    }
    struct Page *page = NULL;
    unsigned index = 0;
    unsigned node_size;
    unsigned offset = 0;
    if (n <= 0)n = 1;
    else if (!IS_POWER_OF_2(n))
    {
        n = 1<<ROUND_DOWN_LOG(n);
    }
    if (buddy[index] < n)
        offset = -1;
    for (node_size = buddy_size; node_size != n; node_size /= 2)
    {
        unsigned left = buddy[LEFT_LEAF(index)];
        unsigned right = buddy[RIGHT_LEAF(index)];

        if (left > right) {
            if (right >= n)
                index = RIGHT_LEAF(index);
            else
                index = LEFT_LEAF(index);
        } else {
            if (left >= n)
                index = LEFT_LEAF(index);
            else
                index = RIGHT_LEAF(index);
        }
    }
    buddy[index] = 0;
    offset = (index + 1) * node_size - buddy_size;
    while (index)
    {
       index=PARENT(index);
        buddy[index]=Max(buddy[LEFT_LEAF(index)],buddy[RIGHT_LEAF(index)]);
    }
    page = base_page+ offset;
    page->property = n;
    unsigned size_free = 1<<ROUND_DOWN_LOG(n);
    nr_free -= size_free;
    for (struct Page *p = page; p != page + size_free; p++)
        ClearPageProperty(p);
    return page;
}

static void
buddy_free_pages(struct Page *base, size_t n)
{
    assert(n > 0);
    n=1<<ROUND_DOWN_LOG(n);
    struct Page *p = base;
    for (; p != base + n; p++)
    {
        assert(!PageReserved(p) && !PageProperty(p));
        set_page_ref(p, 0);
    }
    unsigned offset = base - base_page;
    unsigned node_size = 1;
    unsigned index = buddy_size + offset - 1;
    unsigned left_len, right_len;
    for (; buddy[index]; index = PARENT(index)) {
        node_size *= 2;
        if (!index)
            return;
    }
    buddy[index] = node_size;
    while (index)
    {
        index=PARENT(index);
        node_size *= 2;

        left_len = buddy[LEFT_LEAF(index)];
        right_len = buddy[RIGHT_LEAF(index)];

        if (left_len + right_len == node_size)
            buddy[index] = node_size;
        else
            buddy[index] = Max(left_len, right_len);    
            }

    for (struct Page *p = base; p != base + n; p++)
        ClearPageProperty(p);

    unsigned size_free = 1<<ROUND_DOWN_LOG(n);
    nr_free += size_free;
}

static size_t
buddy_nr_free_pages(void)
{
    return nr_free;
}
//下面检测来自伙伴分配器链接开端那张图+csdn
static void
buddy_check(void)
{
    struct Page *p0, *A, *B, *C, *D;
    p0 = A = B = C = D = NULL;

    assert((p0 = alloc_page()) != NULL);
    assert((A = alloc_page()) != NULL);
    assert((B = alloc_page()) != NULL);

    // assert(p0 != A && p0 != B && A != B);
    // assert(page_ref(p0) == 0 && page_ref(A) == 0 && page_ref(B) == 0);
    cprintf("p0 %p\n", p0);
    cprintf("A %p\n", A);
    cprintf("B %p\n", B);
    free_page(p0);
    free_page(A);
    free_page(B);

//q
    A = alloc_pages(500);
    B = alloc_pages(500);
    cprintf("A %p\n", A);
    cprintf("B %p\n", B);
    free_pages(A, 250);
    free_pages(B, 500);
    free_pages(A + 250, 250);

    p0 = alloc_pages(8192);
    cprintf("p0 %p\n", p0);
    
    //链接中的样例进行测试
    A = alloc_pages(70);
    B = alloc_pages(35);
    cprintf("A t %p\n", A);
    cprintf("B t %p\n", B);
    C = alloc_pages(80);
    cprintf("C %p\n", C);
    free_pages(A, 70);  
    D = alloc_pages(60);
    cprintf("D %p\n", D);
    free_pages(B, 35);
    cprintf("D %p\n", D);
    free_pages(D, 60);
    cprintf("C %p\n", C);
    free_pages(C, 80);
    free_pages(p0, 8192);  
}
const struct pmm_manager buddy_pmm_manager = {
    .name = "buddy_pmm_manager",
    .init = buddy_init,
    .init_memmap = buddy_init_memmap,
    .alloc_pages = buddy_alloc_pages,
    .free_pages = buddy_free_pages,
    .nr_free_pages = buddy_nr_free_pages,
    .check = buddy_check,
};
